import taichi as ti


@ti.data_oriented
class Semi_Lagrangian:
    def __init__(self, grid_n, nParticles, dt):
        ti.init(arch=ti.cpu)
        self.grid_n = grid_n
        self.dt = dt
        self.g = 9.8
        self.pressure = ti.Vector.field(
            2, dtype=ti.f32, shape=(grid_n, grid_n))
        self.velocity_u = ti.field(dtype=ti.f32, shape=(grid_n+1, grid_n))
        self.velocity_v = ti.field(dtype=ti.f32, shape=(grid_n, grid_n+1))
        self.new_velocity_u = ti.field(dtype=ti.f32, shape=(grid_n+1, grid_n))
        self.new_velocity_v = ti.field(dtype=ti.f32, shape=(grid_n, grid_n+1))
        self.grid_length = 1/self.grid_n
        self.particles = ti.Vector.field(2, dtype=ti.f32, shape=(nParticles))
        self.initialize() 
        pass

    @ti.kernel
    def initialize(self):
        for i, j in ti.grouped(self.pressure):
            self.pressure[i, j] = ti.Vector([0.0, 0.0])
        for i, j in self.velocity_u:
            self.velocity_u[i, j] = 0.0
        for i, j in self.velocity_v:
            self.velocity_v[i, j] = 0.0
        for i in self.particles:
            self.particles[i] = ti.Vector([0.0, 0.0])

    def BilinearInterpolation(self):
        pass

    @ti.func
    def Velocity(self, X: ti.template()):
        # 像素点X在第
        base = (X // self.grid_length).cast(int)
        offset = X-base
        Vx = offset[0]*self.velocity_x[base+1]+(1-offset[0])*self.velocity_x[base]
        Vy = offset[1]*self.velocity_y[base+1]+(1-offset[1])*self.velocity_y[base]
        return ti.Vector([Vx,Vy])

    @ti.func
    def advect(self):
        for i,j in self.velocity_u:
            XG = ti.Vector([i,j+0.5])
            XG -= ti.Vector([self.velocity_u[i,j] * self.dt,0])
            self.new_velocity_u += self.Velocity(XG)[0]
        for i,j in self.velocity_v:
            XG = ti.Vector([i+0.5,j])
            XG -= ti.Vector([0,self.velocity_v[i,j] * self.dt])
            self.new_velocity_v += self.Velocity(XG)[1]


    def project(self):
        pass

    def body_force(self):
        for i,j in self.new_velocity_v:
            self.new_velocity_v[i,j] -= self.dt * self.g

    
    @ti.kernel
    def renderGUI(self):
        pass

    @ti.kernel
    def test(self):
        self.advect()
        pass



if __name__ == "__main__":
    semi = Semi_Lagrangian(grid_n=128, nParticles=10000, dt=0.0001)
    semi.test()
